# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os.path
import os.path as osp

import mmcv
import numpy as np

from ..builder import PIPELINES
import json
import torch
from PIL import Image


@PIPELINES.register_module()
class LoadImageFromFile(object):
    """Load an image from file.

    Required keys are "img_prefix" and "img_info" (a dict that must contain the
    key "filename"). Added or updated keys are "filename", "img", "img_shape",
    "ori_shape" (same as `img_shape`), "pad_shape" (same as `img_shape`),
    "scale_factor" (1.0) and "img_norm_cfg" (means=0 and stds=1).

    Args:
        to_float32 (bool): Whether to convert the loaded image to a float32
            numpy array. If set to False, the loaded image is an uint8 array.
            Defaults to False.
        color_type (str): The flag argument for :func:`mmcv.imfrombytes`.
            Defaults to 'color'.
        file_client_args (dict): Arguments to instantiate a FileClient.
            See :class:`mmcv.fileio.FileClient` for details.
            Defaults to ``dict(backend='disk')``.
        imdecode_backend (str): Backend for :func:`mmcv.imdecode`. Default:
            'cv2'
    """

    def __init__(self,
                 to_float32=False,
                 color_type='color',
                 file_client_args=dict(backend='disk'),
                 imdecode_backend='cv2'):
        self.to_float32 = to_float32
        self.color_type = color_type
        self.file_client_args = file_client_args.copy()
        self.file_client = None
        self.imdecode_backend = imdecode_backend

    def __call__(self, results):
        """Call functions to load image and get image meta information.

        Args:
            results (dict): Result dict from :obj:`mmseg.CustomDataset`.

        Returns:
            dict: The dict contains loaded image and meta information.
        """

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)

        if results.get('img_prefix') is not None:
            filename = osp.join(results['img_prefix'],
                                results['img_info']['filename'])
        else:
            filename = results['img_info']['filename']
        img_bytes = self.file_client.get(filename)
        img = mmcv.imfrombytes(
            img_bytes, flag=self.color_type, backend=self.imdecode_backend)
        if self.to_float32:
            img = img.astype(np.float32)

        results['filename'] = filename
        results['ori_filename'] = results['img_info']['filename']
        results['img'] = img
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        # Set initial values for default meta_keys
        results['pad_shape'] = img.shape
        results['scale_factor'] = 1.0
        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
        results['img_norm_cfg'] = dict(
            mean=np.zeros(num_channels, dtype=np.float32),
            std=np.ones(num_channels, dtype=np.float32),
            to_rgb=False)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(to_float32={self.to_float32},'
        repr_str += f"color_type='{self.color_type}',"
        repr_str += f"imdecode_backend='{self.imdecode_backend}')"
        return repr_str


@PIPELINES.register_module()
class LoadAnnotations(object):
    """Load annotations for semantic segmentation.

    Args:
        reduce_zero_label (bool): Whether reduce all label value by 1.
            Usually used for datasets where 0 is background label.
            Default: False.
        file_client_args (dict): Arguments to instantiate a FileClient.
            See :class:`mmcv.fileio.FileClient` for details.
            Defaults to ``dict(backend='disk')``.
        imdecode_backend (str): Backend for :func:`mmcv.imdecode`. Default:
            'pillow'
    """

    def __init__(self,
                 reduce_zero_label=False,
                 file_client_args=dict(backend='disk'),
                 imdecode_backend='pillow'):
        self.reduce_zero_label = reduce_zero_label
        self.file_client_args = file_client_args.copy()
        self.file_client = None
        self.imdecode_backend = imdecode_backend

    def _repeate_op_B(self, results):
        if results.get('B_seg_prefix', None) is not None:
            filename = osp.join(results['B_seg_prefix'],
                                results['B_ann_info']['seg_map'])
        else:
            filename = results['B_ann_info']['seg_map']
        img_bytes = self.file_client.get(filename)
        gt_semantic_seg = mmcv.imfrombytes(
            img_bytes, flag='unchanged',
            backend=self.imdecode_backend).squeeze().astype(np.uint8)
        # modify if custom classes
        if results.get('label_map', None) is not None:
            for old_id, new_id in results['label_map'].items():
                gt_semantic_seg[gt_semantic_seg == old_id] = new_id
        # reduce zero_label
        if self.reduce_zero_label:
            # avoid using underflow conversion
            gt_semantic_seg[gt_semantic_seg == 0] = 255
            gt_semantic_seg = gt_semantic_seg - 1
            gt_semantic_seg[gt_semantic_seg == 254] = 255
        results['B_gt_semantic_seg'] = gt_semantic_seg
        results['seg_fields'].append('B_gt_semantic_seg')

    def __call__(self, results):
        """Call function to load multiple types annotations.

        Args:
            results (dict): Result dict from :obj:`mmseg.CustomDataset`.

        Returns:
            dict: The dict contains loaded semantic segmentation annotations.
        """

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)

        if results.get('seg_prefix', None) is not None:
            filename = osp.join(results['seg_prefix'],
                                results['ann_info']['seg_map'])
        else:
            filename = results['ann_info']['seg_map']
        img_bytes = self.file_client.get(filename)
        gt_semantic_seg = mmcv.imfrombytes(
            img_bytes, flag='unchanged',
            backend=self.imdecode_backend).squeeze().astype(np.uint8)
        # modify if custom classes
        if results.get('label_map', None) is not None:
            for old_id, new_id in results['label_map'].items():
                gt_semantic_seg[gt_semantic_seg == old_id] = new_id
        # reduce zero_label
        if self.reduce_zero_label:
            # avoid using underflow conversion
            gt_semantic_seg[gt_semantic_seg == 0] = 255
            gt_semantic_seg = gt_semantic_seg - 1
            gt_semantic_seg[gt_semantic_seg == 254] = 255
        results['gt_semantic_seg'] = gt_semantic_seg
        results['seg_fields'].append('gt_semantic_seg')

        # print('load info A:\t',results['img'].shape, results['gt_semantic_seg'].shape)

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(reduce_zero_label={self.reduce_zero_label},'
        repr_str += f"imdecode_backend='{self.imdecode_backend}')"
        return repr_str


@PIPELINES.register_module()
class LoadAnnotations_B(object):
    """Load annotations for semantic segmentation.
    please read LoadAnnotations"""

    def __init__(self,
                 reduce_zero_label=False,
                 file_client_args=dict(backend='disk'),
                 imdecode_backend='pillow'):
        self.reduce_zero_label = reduce_zero_label
        self.file_client_args = file_client_args.copy()
        self.file_client = None
        self.imdecode_backend = imdecode_backend

    def __call__(self, results):
        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)

        if results.get('B_seg_prefix', None) is not None:
            filename = osp.join(results['B_seg_prefix'],
                                results['B_ann_info']['seg_map'])
        else:
            filename = results['B_ann_info']['seg_map']

        img_bytes = self.file_client.get(filename)
        gt_semantic_seg = mmcv.imfrombytes(
            img_bytes, flag='unchanged',
            backend=self.imdecode_backend).squeeze().astype(np.uint8)
        # modify if custom classes
        if results.get('label_map', None) is not None:
            for old_id, new_id in results['label_map'].items():
                gt_semantic_seg[gt_semantic_seg == old_id] = new_id
        # reduce zero_label
        if self.reduce_zero_label:
            # avoid using underflow conversion
            gt_semantic_seg[gt_semantic_seg == 0] = 255
            gt_semantic_seg = gt_semantic_seg - 1
            gt_semantic_seg[gt_semantic_seg == 254] = 255
        results['B_gt_semantic_seg'] = gt_semantic_seg
        results['seg_fields'].append('B_gt_semantic_seg')

        # print('load info B:\t',results['B_img'].shape, results['B_gt_semantic_seg'].shape)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(B_reduce_zero_label={self.reduce_zero_label},'
        repr_str += f"B_imdecode_backend='{self.imdecode_backend}')"
        return repr_str


@PIPELINES.register_module()
class LoadAnnotations_B_active(object):
    """Load annotations for semantic segmentation.
    please read LoadAnnotations"""

    def __init__(self,reduce_zero_label=False ):
        self.reduce_zero_label=reduce_zero_label

    def __call__(self, results):
        if 'active_mask_path' in results['B_img_info']:
            partial_mask_path = results['B_img_info']['active_mask_path']
            partial_mask = np.array(Image.open(partial_mask_path)).astype(np.uint8)
        else:
            partial_mask = np.zeros_like(results['B_gt_semantic_seg'])

        if self.reduce_zero_label:
            # avoid using underflow conversion
            partial_mask[partial_mask == 0] = 255
            partial_mask = partial_mask - 1
            partial_mask[partial_mask == 254] = 255
        results['B_partial_semantic_seg'] = partial_mask
        results['seg_fields'].append('B_partial_semantic_seg')
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        return repr_str


## added by LYU: 2022/03/24
@PIPELINES.register_module()
class LoadImageFromFile_forAdap(object):
    """Load an image from file.

    Required keys are "img_prefix" and "img_info" (a dict that must contain the
    key "filename"). Added or updated keys are "filename", "img", "img_shape",
    "ori_shape" (same as `img_shape`), "pad_shape" (same as `img_shape`),
    "scale_factor" (1.0) and "img_norm_cfg" (means=0 and stds=1).

    Args:
        to_float32 (bool): Whether to convert the loaded image to a float32
            numpy array. If set to False, the loaded image is an uint8 array.
            Defaults to False.
        color_type (str): The flag argument for :func:`mmcv.imfrombytes`.
            Defaults to 'color'.
        file_client_args (dict): Arguments to instantiate a FileClient.
            See :class:`mmcv.fileio.FileClient` for details.
            Defaults to ``dict(backend='disk')``.
        imdecode_backend (str): Backend for :func:`mmcv.imdecode`. Default:
            'cv2'
    """

    def __init__(self,
                 to_float32=False,
                 color_type='color',
                 file_client_args=dict(backend='disk'),
                 imdecode_backend='cv2'):
        self.to_float32 = to_float32
        self.color_type = color_type
        self.file_client_args = file_client_args.copy()
        self.file_client = None
        self.imdecode_backend = imdecode_backend

    def __call__(self, results):
        """Call functions to load image and get image meta information.

        Args:
            results (dict): Result dict from :obj:`mmseg.CustomDataset`.

        Returns:
            dict: The dict contains loaded image and meta information.
        """

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)

        ## added by LYU: 2022/03/24 img of domain A
        if results.get('img_prefix') is not None:
            filename = osp.join(results['img_prefix'],
                                results['img_info']['filename'])
        else:
            filename = results['img_info']['filename']
        img_bytes = self.file_client.get(filename)
        img = mmcv.imfrombytes(
            img_bytes, flag=self.color_type, backend=self.imdecode_backend)
        if self.to_float32:
            img = img.astype(np.float32)

        ## added by LYU: 2022/03/24 img of domain B
        if results.get('B_img_prefix') is not None:
            B_filename = osp.join(results['B_img_prefix'],
                                  results['B_img_info']['filename'])
        else:
            B_filename = results['B_img_info']['filename']
        B_img_bytes = self.file_client.get(B_filename)
        B_img = mmcv.imfrombytes(
            B_img_bytes, flag=self.color_type, backend=self.imdecode_backend)
        if self.to_float32:
            B_img = B_img.astype(np.float32)

        ## added by LYU: 2022/03/24 results of domain A
        results['filename'] = filename
        results['ori_filename'] = results['img_info']['filename']
        results['img'] = img
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        # Set initial values for default meta_keys
        results['pad_shape'] = img.shape
        results['scale_factor'] = 1.0
        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
        results['img_norm_cfg'] = dict(
            mean=np.zeros(num_channels, dtype=np.float32),
            std=np.ones(num_channels, dtype=np.float32),
            to_rgb=False)

        ## added by LYU: 2022/03/24 results of domain B
        results['B_filename'] = B_filename
        results['B_ori_filename'] = results['B_img_info']['filename']
        results['B_img'] = B_img
        results['B_img_shape'] = B_img.shape
        results['B_ori_shape'] = B_img.shape
        # Set initial values for default meta_keys
        results['B_pad_shape'] = B_img.shape
        results['B_scale_factor'] = 1.0
        B_num_channels = 1 if len(B_img.shape) < 3 else B_img.shape[2]
        results['B_img_norm_cfg'] = dict(
            mean=np.zeros(B_num_channels, dtype=np.float32),
            std=np.ones(B_num_channels, dtype=np.float32),
            to_rgb=False)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(to_float32={self.to_float32},'
        repr_str += f"color_type='{self.color_type}',"
        repr_str += f"imdecode_backend='{self.imdecode_backend}')"
        return repr_str


from mmseg.utils.sam_tools import RandomMaskSampler, gen_comb_mask


@PIPELINES.register_module()
class LoadCoCoMask_B(object):
    def __init__(self, ratio=None):
        self.ratio = ratio

    def __call__(self, results):
        filename = results['B_img_info']['filename'].split('.')[0] + '.json'
        filepath = os.path.join(results['B_coco_mask_prefix'], filename)
        with open(filepath, 'r') as f:
            masks = json.load(f)
        mask_sampler = RandomMaskSampler(masks)
        ratio = self.ratio if self.ratio is not None else 0.4 + np.random.random() * 0.2
        select_masks = mask_sampler.sample(ratio)
        comb_masks = gen_comb_mask(select_masks)
        results['B_auto_mask'] = comb_masks
        # results['seg_fields'].append('B_auto_mask')
        # print(f'load B_auto_mask:{filename}')

        return results

    def __repr__(self, ):
        repr_str = self.__class__.__name__
        return repr_str
